import warnings
from functools import wraps
from typing import Callable, Optional, Union


class Registry(dict):
    """
    This calss can register or find a callable object by string. Tt like a
    python dict, but the difference is registration way, this class use
    python decorator instead of `dict[key] = value`.

    Args:
        registry_name (str, optional): Specify name for registry to be created.
        description (str, optional): Describe the application scenario for
            registrar to be created.

    Example:

        step1. Create a registry.
        >>> from npu_mmdet.utils import Registry
        >>> TRANSFORMS =  Registry(
        >>>     registry_name='TRANSFORMS',
        >>>     description='Transforms for data preprocessing')

        step2. Register a funciton.
        >>> @TRANSFORMS.register('Normalize')
        >>> def normalize(img, mean=[123.675, 116.28, 103.53],
        >>>               std=[58.395, 57.12, 57.375],):
        >>>     img = (img - np.array(mean)) / np.array(std)
        >>>     return img.astype(np.float32)

        Of course, you can also register explicitly.
        >>> TRANSFORMS.register(normalize, 'Normalize')
    """

    def __init__(self,
                 registry_name: Optional[str] = None,
                 description: Optional[str] = None):
        super(Registry, self).__init__()
        self.name = registry_name
        self.description = description

    def register(self, arg: Union[str, Callable]) -> Callable:
        def add_item(target: Callable, name: Optional[str] = None):
            assert callable(target), 'target must be a function or class!'
            if not name:
                name = target.__name__
            if name in self:
                warnings.warn(
                    f'{name} has been registered:\n'
                    f'\tkey = {name}\n'
                    f'\ttarget = {self[name]}\n'
                    'covering target may cause changes in the prediction '
                    'results of model.',
                    category=RuntimeWarning,
                    source=None
                )
            self[name] = target
            return target

        if callable(arg):
            return add_item(arg)

        return lambda target: add_item(target, name=arg)

    def __call__(self, *args, **kwargs):
        return self.register(*args, **kwargs)
